//===----------------------------------------------------------------------===//
//
//                         BusTub
//
// hash_join_executor.cpp
//
// Identification: src/execution/hash_join_executor.cpp
//
// Copyright (c) 2015-2021, Carnegie Mellon University Database Group
//
//===----------------------------------------------------------------------===//

#include "execution/executors/hash_join_executor.h"

namespace bustub {

HashJoinExecutor::HashJoinExecutor(ExecutorContext *exec_ctx, const HashJoinPlanNode *plan,
                                   std::unique_ptr<AbstractExecutor> &&left_child,
                                   std::unique_ptr<AbstractExecutor> &&right_child)
    : AbstractExecutor(exec_ctx),
      plan_(plan),
      left_child_(std::move(left_child)),
      right_child_(std::move(right_child)) {
  if (!(plan->GetJoinType() == JoinType::LEFT || plan->GetJoinType() == JoinType::INNER)) {
    // Note for 2023 Spring: You ONLY need to implement left join and inner join.
    throw bustub::NotImplementedException(fmt::format("join type {} not supported", plan->GetJoinType()));
  }
}

void HashJoinExecutor::Init() {
  left_child_->Init();
  right_child_->Init();
  Tuple tuple{};
  RID rid{};
  while (left_child_->Next(&tuple, &rid)) {
    AddTupleToHashTable(tuple, left_hash_table_, plan_->LeftJoinKeyExpressions());
  }
  while (right_child_->Next(&tuple, &rid)) {
    AddTupleToHashTable(tuple, right_hash_table_, plan_->RightJoinKeyExpressions());
  }
  left_hash_table_iterator_ = left_hash_table_.begin();
  if (left_hash_table_iterator_ != left_hash_table_.end()) {
    left_tuple_iterator_ = left_hash_table_iterator_->second.tuples_.begin();
  }
}

auto HashJoinExecutor::Next(Tuple *tuple, RID *rid) -> bool {
  while (left_hash_table_iterator_ != left_hash_table_.end()) {
    if (!right_initialized_ && right_hash_table_.count(left_hash_table_iterator_->first) != 0) {
      right_hash_table_iterator_ = right_hash_table_.find(left_hash_table_iterator_->first);
      right_tuple_iterator_ = right_hash_table_iterator_->second.tuples_.begin();
      right_initialized_ = true;
    }
    while (left_tuple_iterator_ != left_hash_table_iterator_->second.tuples_.end()) {
      while (right_initialized_ && right_tuple_iterator_ != right_hash_table_iterator_->second.tuples_.end()) {
        const Tuple &left_tuple = *left_tuple_iterator_;
        const Tuple &right_tuple = *right_tuple_iterator_;
        JoinTuple(left_tuple, right_tuple, tuple);
        ++right_tuple_iterator_;
        return true;
      }
      if (plan_->GetJoinType() == JoinType::LEFT && !right_initialized_) {
        const Tuple &left_tuple = *left_tuple_iterator_;
        std::vector<bustub::Value> vals{};
        for (decltype(plan_->GetLeftPlan()->OutputSchema().GetColumnCount()) i = 0;
             i < plan_->GetLeftPlan()->OutputSchema().GetColumnCount(); i++) {
          vals.emplace_back(left_tuple.GetValue(&plan_->GetLeftPlan()->OutputSchema(), i));
        }
        for (decltype(plan_->GetRightPlan()->OutputSchema().GetColumnCount()) i = 0;
             i < plan_->GetRightPlan()->OutputSchema().GetColumnCount(); i++) {
          vals.emplace_back(
              ValueFactory::GetNullValueByType(plan_->GetRightPlan()->OutputSchema().GetColumn(i).GetType()));
        }
        *tuple = Tuple{vals, &plan_->OutputSchema()};
        ++left_tuple_iterator_;
        return true;
      }
      if (right_initialized_) {
        right_tuple_iterator_ = right_hash_table_iterator_->second.tuples_.begin();
      }
      ++left_tuple_iterator_;
    }
    ++left_hash_table_iterator_;
    if (left_hash_table_iterator_ != left_hash_table_.end()) {
      left_tuple_iterator_ = left_hash_table_iterator_->second.tuples_.begin();
    }
    right_initialized_ = false;
  }
  return false;
}

void HashJoinExecutor::AddTupleToHashTable(const Tuple &tuple,
                                           std::unordered_map<AggregateKey, HashJoinTuples> &hash_table,
                                           const std::vector<bustub::AbstractExpressionRef> &join_key_expressions) {
  AggregateKey k{};
  for (const auto &join_key_expression : join_key_expressions) {
    const auto &value = join_key_expression->Evaluate(&tuple, plan_->GetLeftPlan()->OutputSchema());
    k.group_bys_.emplace_back(value);
  }
  if (hash_table.count(k) == 0) {
    HashJoinTuples tuples{};
    tuples.tuples_.emplace_back(tuple);
    hash_table[k] = tuples;
  } else {
    hash_table[k].tuples_.emplace_back(tuple);
  }
}

auto HashJoinExecutor::JoinTuple(const Tuple &left_tuple, const Tuple &right_tuple, Tuple *tuple) -> bool {
  std::vector<bustub::Value> vals{};
  for (decltype(plan_->GetLeftPlan()->OutputSchema().GetColumnCount()) i = 0;
       i < plan_->GetLeftPlan()->OutputSchema().GetColumnCount(); i++) {
    vals.emplace_back(left_tuple.GetValue(&plan_->GetLeftPlan()->OutputSchema(), i));
  }
  for (decltype(plan_->GetRightPlan()->OutputSchema().GetColumnCount()) i = 0;
       i < plan_->GetRightPlan()->OutputSchema().GetColumnCount(); i++) {
    vals.emplace_back(right_tuple.GetValue(&plan_->GetRightPlan()->OutputSchema(), i));
  }
  *tuple = Tuple{vals, &plan_->OutputSchema()};
  return true;
}

}  // namespace bustub
